# -*- coding: utf-8 -*-
# @Time    : 2023/6/27
# @Author  : 电信2006赵明涛
# @File    : main.py
# @Software: Pycharm

import os
import sys
import cv2
import numpy as np
from matplotlib import pyplot as plt
import functools
from tkinter import Tk, Button
from tkinter.filedialog import askopenfilename
from tkinter.constants import *


class IMG_AFFINE(object):
    def __init__(self, filepath):
        self.imgpath = filepath
        self.img = cv2.imread(filepath)
        plt.figure()
        plt.imshow(self.img)
        plt.title("original")
        self.gray = cv2.cvtColor(self.img, cv2.COLOR_BGR2GRAY)
        plt.figure()
        plt.imshow(self.gray, cmap="gray")
        plt.title("gray")
        # 二值化处理,使用OTSU算法自动计算阈值,cv2.THRESH_BINARY_INV黑白值反色
        ret, self.binary = cv2.threshold(self.gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        plt.figure()
        plt.imshow(self.binary, cmap="gray")
        plt.title("binary")

    # Morphological treatment形态学处理
    # 闭运算
    def mt(self):
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
        # 使用的核为 kernel，迭代次数为1
        self.closed = cv2.morphologyEx(self.binary, cv2.MORPH_CLOSE, kernel, iterations=1)
        plt.figure()
        plt.imshow(self.closed, cmap="gray")
        plt.title("closed")

    def line_detection(self):
        # rho：以像素为单位的距离精度。theta：以弧度为单位的角度精度。threshold：阈值参数。minLineLength：线段最小长度。比这个短的线段将被拒绝。maxLineGap：线段之间允许的最大间隙。
        lines = cv2.HoughLinesP(self.closed, 1, np.pi / 180, threshold=80, minLineLength=100, maxLineGap=10)
        # 找出直线的平均角度
        angles = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            angle = np.arctan2(y2 - y1, x2 - x1) * 180.0 / np.pi
            angles.append(angle)
        self.mean_angle = np.mean(angles)

    def img_trans(self):
        # 获取原图长宽
        h, w = self.img.shape[:2]
        # 计算旋转矩阵，缩放比例不变
        # M = [
        #     [cos(a), -sin(a), (1 - cos(a)) * w / 2 + sin(a) * h / 2],
        #     [sin(a), cos(a), (1 - cos(a)) * h / 2 - sin(a) * w / 2]
        # ]
        M = cv2.getRotationMatrix2D((w / 2, h / 2), self.mean_angle, 1)
        center = (w // 2, h // 2)
        cos = np.abs(M[0, 0])
        sin = np.abs(M[0, 1])
        # 新图的长宽
        nw = int((h * sin) + (w * cos))
        nh = int((h * cos) + (w * sin))
        # 新图像的中心
        M[0, 2] += (nw / 2) - center[0]
        M[1, 2] += (nh / 2) - center[1]

        self.rotated = cv2.warpAffine(self.img, M, (nw, nh))
        plt.figure()
        plt.imshow(self.rotated)
        plt.title("result")

    def get_img(self):
        self.mt()
        self.line_detection()
        self.img_trans()
        plt.show()
        cv2.imwrite('result.jpg', self.rotated)


class IMG_PERSPECTIVE(object):
    points = []

    def __init__(self, filepath):
        self.img = cv2.imread(filepath)
        cv2.namedWindow('image', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('image', self.img.shape[1], self.img.shape[0])

        callback_func = functools.partial(self.on_mouse)
        cv2.setMouseCallback('image', callback_func)

    def on_mouse(self, event, x, y, flags, param):
        # 处理鼠标事件，并修改图像
        if event == cv2.EVENT_LBUTTONUP:
            self.points.append((x, y))
            cv2.circle(self.img, (x, y), 2, (0, 0, 255), thickness=-1)
            # cv2.imshow('image', self.img)

    def get_img(self):
        while True:
            cv2.imshow('image', self.img)
            key = cv2.waitKey(1)
            if key == ord('q'):
                break
            elif key == ord('c'):
                if len(self.points) == 4:
                    # 计算透视变换矩阵
                    pts1 = np.float32(self.points)
                    width = max(
                        np.linalg.norm(pts1[0] - pts1[1]),
                        np.linalg.norm(pts1[2] - pts1[3])
                    )
                    height = max(
                        np.linalg.norm(pts1[0] - pts1[2]),
                        np.linalg.norm(pts1[1] - pts1[3])
                    )
                    pts2 = np.float32([[0, 0], [width, 0], [0, height], [width, height]])
                    # 构造透视变换矩阵
                    M = cv2.getPerspectiveTransform(pts1, pts2)

                    # 应用透视变换
                    result = cv2.warpPerspective(self.img, M, (int(width), int(height)))  # 传递计算的输出图像大小

                    # 显示结果
                    plt.imshow(result)
                    plt.title("result")
                    plt.show()
                    cv2.imwrite('result.jpg', result)
                    break


def img_fun1(filepath):
    obj = IMG_PERSPECTIVE(filepath)
    obj.get_img()
    sys.exit()


def img_fun2(filepath):
    obj = IMG_AFFINE(filepath)
    obj.get_img()
    sys.exit()


if __name__ == '__main__':
    current_dir = os.getcwd()
    # 创建根窗口，隐藏 GUI 界面
    root = Tk()
    root.withdraw()
    # 打开文件选择对话框
    file_path = askopenfilename(
        initialdir=current_dir)

    img = cv2.imread(file_path)
    plt.figure()
    plt.imshow(img)
    plt.title("original")
    plt.show()

    root = Tk()
    root.title("选择")

    button1 = Button(root, text="存在垂直旋转", command=lambda: img_fun1(file_path))
    button2 = Button(root, text="仅存在水平旋转", command=lambda: img_fun2(file_path))

    button1.pack(side=LEFT, padx=10, pady=10)
    button2.pack(side=RIGHT, padx=10, pady=10)

    root.geometry("400x300+100+100")
    root.mainloop()
